#include "ArithmeticDecoder.h"
#include <string.h>
#include <iostream>

ArithmeticDecoder::ArithmeticDecoder(ProbabilityModel* model, unsigned int outputSize)
{
	this->model = model;
	lowerBound = 0;
	upperBound = 0xFFFFFFFF;
	underflow = 0;
	code = 0;
	result = new char[outputSize];
	ctwTree = NULL;
	for(int i = 0; i < outputSize; i++)
	{
		result[i] = 0;
	}
	symbolSize = model->getSymbolSize();
	bitsInRegister = 0;
}

ArithmeticDecoder::ArithmeticDecoder(ContextTree* tree, int modelSize, int blockSize, unsigned int outputSize)
{
	this->ctwTree = tree;
	this->modelSize = modelSize;
	this->blockSize = blockSize;
	model = NULL;
	lowerBound = 0;
	upperBound = 0xFFFFFFFF;
	underflow = 0;
	code = 0;
	result = new char[outputSize];
	for(int i = 0; i < outputSize; i++)
	{
		result[i] = 0;
	}
	bitsInRegister = 0;
	symbolSize = 1;
}


ArithmeticDecoder::~ArithmeticDecoder(void)
{
	delete result;
}


char* ArithmeticDecoder::decodeBitStream(BitBuffer& buffer, unsigned long long symbolsToDecode)
{
	unsigned char bitmask = 0x80;
	unsigned char shift = 7;
	int i = 0;
	int symbolsInTree = 0;
	// init first characters in tree
	if (ctwTree != NULL)
	{
		i += modelSize;
		BitBuffer treeBuffer(modelSize, 1);
		for (int j = 0; j < modelSize; ++j)
		{
			char bit = buffer.popBit();
			result[j / 8] |= (bit << shift) & bitmask;
			treeBuffer.appendBit(bit);
			--shift;
			bitmask >>= 1;
			if (bitmask == 0)
			{
				bitmask = 0x80;
				shift = 7;
			}
		}
		treeBuffer.resetPointers();
		ctwTree->countTree(&treeBuffer, modelSize);
		symbolsInTree = modelSize;
		symbolsToDecode -= modelSize;
	}

	this->initializeDecoder(buffer);
	for (;;)
	{
		// tree full
		if (ctwTree != NULL && symbolsInTree % (blockSize * 8) == 0)
		{
			lowerBound = 0;
			upperBound = 0xFFFFFFFF;
			underflow = 0;
			code <<= 1;
			if (buffer.hasBits())
			{
				code |= buffer.popBit();
			}
			code <<= 1;
			if (buffer.hasBits())
			{
				code |= buffer.popBit();
			}

			BitBuffer treeBuffer(modelSize, 1);
			for (int j = 0; j < modelSize; ++j, ++i)
			{
				char bit = (code & 0x80000000) >> 31;
				result[i / 8] |= (bit << shift) & bitmask;
				treeBuffer.appendBit(bit);
				--shift;
				bitmask >>= 1;
				code <<= 1;
				if (buffer.hasBits())
				{
					code |= buffer.popBit();
				}
				if (bitmask == 0)
				{
					bitmask = 0x80;
					shift = 7;
				}
				--symbolsToDecode;
				if (symbolsToDecode == 0)
				{
					break;
				}
			}
			if (symbolsToDecode == 0)
			{
				break;
			}
			treeBuffer.resetPointers();
			ctwTree->countTree(&treeBuffer, modelSize);
			symbolsInTree = modelSize;
		}
		unsigned int symbol = decodeSymbol(buffer);
		--symbolsToDecode;
		++symbolsInTree;
		for (unsigned int j = 0; j < symbolSize; ++j, ++i)
		{
			result[i / 8] |= ((symbol >> (symbolSize - j - 1)) << shift) & bitmask;
			bitmask >>= 1;
			--shift;
			if (bitmask == 0)
			{
				bitmask = 0x80;
				shift = 7;
			}
		}
		if (symbolsToDecode == 0)
		{
			break;
		}
		this->normalizeRegisters(symbol, buffer);
	}

	return result;
}

unsigned int ArithmeticDecoder::decodeSymbol(BitBuffer& buffer)
{
	long long range = (long long)upperBound - lowerBound + 1;
	unsigned int codeNumber = 0;
	if (symbolSize == 1)
	{
		unsigned int R1;
		if (ctwTree)
		{
			R1 = (unsigned long long)range * ctwTree->predict(0);
		}
		else
		{
			R1 = (unsigned long long)range * model->getUpperBoundForSymbol(0) / model->getAllSymbolsCount();
		}
		codeNumber = code - lowerBound;
		if (codeNumber < R1)
		{
			upperBound = lowerBound + R1 - 1;
			if (ctwTree != NULL)
			{
				ctwTree->appendBitToTree(0);
			}
			return 0;
		}
		else
		{
			lowerBound = lowerBound + R1;
			if (ctwTree != NULL)
			{
				ctwTree->appendBitToTree(1);
			}
			
			return 1;
		}
	}
	else
	{
		codeNumber = ((code - lowerBound + 1) * (unsigned long long)model->getAllSymbolsCount() - 1) / range;
	}
	return model->getSymbolFromCodeNumber(codeNumber);
}

void ArithmeticDecoder::normalizeRegisters(unsigned int symbol, BitBuffer& buffer)
{
	long long range = (long long)upperBound - lowerBound + 1;
	if (symbolSize != 1)
	{
		upperBound = lowerBound + ((unsigned long long)range * model->getUpperBoundForSymbol(symbol) / model->getAllSymbolsCount()) - 1;
		lowerBound = lowerBound + ((unsigned long long)range * model->getLowerBoundForSymbol(symbol) / model->getAllSymbolsCount());
	}
	
	while(true) 
	{
		if ((upperBound & 0x80000000) == (lowerBound & 0x80000000))
		{
		} 
		else if (((upperBound & 0x40000000) == 0) && (lowerBound & 0x40000000)) 
		{
			code ^= 0x40000000;
			upperBound |= 0x40000000;
			lowerBound &= 0x3fffffff;
		}
		else 
		{
			return;
		}

		upperBound <<= 1;
		upperBound |= 1;
		lowerBound <<= 1;
		code <<= 1;
		if (buffer.hasBits())
		{
			code |= buffer.popBit();
		}
		else
		{
			--bitsInRegister;
		}
		if (bitsInRegister == 0)
		{
			return;
		}
	}
}


void ArithmeticDecoder::initializeDecoder(BitBuffer& buffer)
{
	for (int i = 0; i < sizeof(code) * 8; ++i) 
	{
		code <<= 1;
		code |= buffer.popBit();
		++bitsInRegister;
	}
}